package me.zhengjie.controller;

import me.zhengjie.annotation.AnonymousAccess;
import me.zhengjie.common.DlaServerConfig;
import me.zhengjie.entity.QueryRecord;
import me.zhengjie.service.QueryService;
import me.zhengjie.util.MapUtil;
import com.fasterxml.jackson.databind.ObjectMapper;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.Response;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RestController;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;

import static me.zhengjie.util.HttpRequestUtil.getRequiredParameter;
import static io.prestosql.client.OkHttpUtil.basicAuth;
import static javax.servlet.http.HttpServletResponse.SC_OK;

@RestController
public class QueryController {
    private static final int LIMIT = 100;
    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

    @Resource
    private DlaServerConfig config;
    @Resource
    private QueryService queryService;

    private final OkHttpClient httpClient = new OkHttpClient();


    @PostMapping("query")
    @AnonymousAccess
    public void query( HttpServletRequest request, HttpServletResponse response) throws IOException {
        response.setContentType("application/json");
        PrintWriter writer = response.getWriter();
        String dataSource = getRequiredParameter(request, "dataSource");
        String coordinatorServer = config.getPrestoCoordinatorServerOrNull(dataSource);
        if (coordinatorServer == null) {
            writer.println("[]");
            return;
        }

        Optional<String> user = Optional.ofNullable(request.getParameter("user"));
        Optional<String> password = Optional.ofNullable(request.getParameter("password"));
        OkHttpClient client = httpClient;
        if (user.isPresent() && password.isPresent()) {
            if (user.get().isEmpty()) {
                writer.println(OBJECT_MAPPER.writeValueAsString(MapUtil.of("error", "user is empty")));
                return;
            }
            OkHttpClient.Builder clientBuilder = httpClient.newBuilder();
            clientBuilder.addInterceptor(basicAuth(user.get(), password.get()));
            client = clientBuilder.build();
        }

        Request prestoRequest = new Request.Builder().url(coordinatorServer + "/v1/query").build();
        String originalJson;
        try (Response prestoResponse = client.newCall(prestoRequest).execute()) {
            originalJson = prestoResponse.body().string();
            int code = prestoResponse.code();
            if (code != SC_OK) {
                writer.println(OBJECT_MAPPER.writeValueAsString(MapUtil.of("code", code, "error", prestoResponse.message())));
                return;
            }
        }
        List<Map> list = OBJECT_MAPPER.readValue(originalJson, List.class);
        List<Map> runningList = list.stream().filter(m -> m.get("state").equals("RUNNING")).collect(Collectors.toList());
        List<Map> notRunningList = list.stream().filter(m -> !m.get("state").equals("RUNNING")).collect(Collectors.toList());
        runningList.sort((a, b) -> String.class.cast(b.get("queryId")).compareTo(String.class.cast(a.get("queryId"))));
        notRunningList.sort((a, b) -> String.class.cast(b.get("queryId")).compareTo(String.class.cast(a.get("queryId"))));
        List<Map> limitedList = new ArrayList<>();
        limitedList.addAll(runningList);
        limitedList.addAll(notRunningList.subList(0, Math.min(LIMIT, list.size()) - runningList.size()));
        List<String> queryIds = limitedList.stream().map(query -> (String) query.get("queryId")).collect(Collectors.toList());
        Map params = MapUtil.of("engine", "presto", "DataSource", dataSource, "queryIds", queryIds);
        List<QueryRecord> queries = queryService.selectAllByEngineAndDatasourceAndQueryIdIn(params);
        List<String> existDbQueryIds = new ArrayList<>();
        for (QueryRecord queryRecord : queries) {
            existDbQueryIds.add(queryRecord.getQueryId());
        }
        for (Map query : limitedList) {
            String queryId = (String) query.get("queryId");
            query.put("existdb", existDbQueryIds.contains(queryId));
        }
        String json = OBJECT_MAPPER.writeValueAsString(limitedList);
        writer.println(json);
    }
}
